#!/bin/bash
#
# w11arm_esd2iso - download and convert Microsoft ESD files for Windows 11 ARM to ISO
#
# Copyright (C) 2023 Paul Rockwell
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA#
#
# Credit: Location and methods of obtaining Microsoft ESD distributions and
# Microsoft Product catalog from b0gdanw "ESD to ISO on macOS.txt" https://gist.github.com/b0gdanw/e36ea84828dbd19e03eff6158f1fc77c
#


readonly versionID="v5.0.3 (20240221)"
readonly version="w11arm_esd2iso ${versionID}\n"
readonly minDiskSpace=12
readonly requires="wimlib-imagex aria2c cabextract mkisofs xpath shasum"
readonly keepDownloads=0
declare -a lTags
declare -a lDesc
declare esdURL
declare esdFile
declare buildName
declare verbosityLevel=0
declare hOption=0
declare rOption=0
declare dOption=0
declare shaHash=""
declare runningOnDarwin=0
declare tempDir="$(pwd)"

usage() {
	echo -e "Usage:\n"
	echo -e "$0 [-v]"
	echo -e "$0 [-v] -r work-dir"
	echo -e "$0 [-Vh]"
	echo -e "\nOptions:"
	echo -e "\t-h\tPrint usage and exit"
	echo -e "\t-v\tEnable verbose output"
	echo -e "\t-V\tPrint program version and exit"
	echo -e "\t-r work-dir\n\t\tRestart an interrupted ESD download from a prior execution using work-dir"
	echo
}

abortProg() {
	local -r errCode=$1
	local -r errRoutine=$2
	local -r errInfo=$3

	echo -e "[ERROR] $errRoutine failed with error code $errCode"
	echo -e "$errInfo"
	echo -e "Program aborting"
	exit 1
}

verboseOn() {
	if (( verbosityLevel == 0 )); then
		return 1
	else
		return 0
	fi
}

debugOn() {
	if (( dOption == 0 )) ; then
		return 1
	else
		return 0
	fi
}

printLanguages() {
	local i
	
	printf '\n%s\n\n' "Available Windows 11 ARM languages are:"
	printf '%-15s %-20s\n' "Language tag"	"Language"
	printf '%-15s %-20s\n' "------------" "--------------------"
    for (( i=1; i<=${#lTags[@]}; i++ )); do
    	printf '%-15s %-20s\n' ${lTags[$i]} "${lDesc[$i]}"
    done
    return 0
}
downloadCatalog() {
	local -r winCatalog="https://go.microsoft.com/fwlink?linkid=2156292"
	
	[[ -d ${wDir} ]] ||
		abortProg 2 "Catalog download" "${wDir} is expected to exist but doesn't\nPlease report this error"
	
	#---------------
	# Download and process the Windows Product Catalog so we know what languages are available
	# and the URLs for the downloads
	#---------------
	

	aria2c --dir="${wDir}" --out=catalog.cab --download-result=hide --console-log-level=error "${winCatalog}"  || return $?

	(
		cd "${wDir}"; 
		cabextract catalog.cab > cabextract.log;
		verboseOn && cat cabextract.log 
		sed -i.bak 's/\r//' products.xml 
	)
	return 0
}

setupEsdDownload(){
   	
   	local edQuery
	local editionName=""
	local answer
	local pDesc
	local ltg
	local found
   
    #
	# Ask the user for the ISO edition
	#

	echo -e "\nWhich Edition of Windows 11 ARM do you wish to create?"

	while [[ -z ${editionName} ]]; do

		echo -e "Enter 'p' to create a Professional Edition ISO"
		echo -e "          (includes Professional and Home Editions), or"
		echo -e "      'e' to create a Enterprise Edition ISO"
		echo -e "          (includes Professional and Enterprise Editions), or"
		echo -e "      'q' to exit the program without creating an ISO: \c"
		read answer

		case ${answer} in
			p | P )
				editionName="Professional"
				;;
			e | E )
				editionName="Enterprise"
				;;
			q | Q )
				return 1
				;;
			* )
				echo -e "Invalid response. Please try again.\n"
				;;
		esac
	done

    echo -e "\nDetermining available languages for ${editionName} edition.\c"
    #
	# Get the W11 ARM64 products for the desired edition from the catalog
	#
	
    edQuery='//File[Architecture="ARM64"][Edition="'${editionName}'"]'
	echo -e '<Catalog>' > "${wDir}/w11arm_products.xml"
	echo -e ".\c"
	xpath -q -n -e ${edQuery} "${wDir}/products.xml" >> "${wDir}/w11arm_products.xml"
	echo -e '</Catalog>'>> "${wDir}/w11arm_products.xml"
	echo -e ".\c"
	#
	# Load the available languages so the user can select one
	# 
	
    for ltg in $(xpath -q -n -e '//File/LanguageCode' "${wDir}/w11arm_products.xml" | sed -E -e 's/<[\/]?LanguageCode>//g' | sort); do
    	lTags+=("${ltg}")
    	pDesc=$(xpath -q -n -e '//File[LanguageCode="'$ltg'"]/Language' "${wDir}/w11arm_products.xml" | sed -E -e "s/<[\/]?Language>//g" )
    	lDesc+=("$pDesc")
    	echo -e ".\c"
    done
    
    echo -e "done."
    #
	# Ask the user for the ISO language they want 
	#
	
	echo -e "\n\nPlease select a language for your Windows 11 ARM ${editionName} ISO"
# 
	esdLang=""
	while [[ -z ${esdLang} ]]; do
		echo -e "Enter a Windows language tag (example: 'en-us'), or"
		echo -e "     'p' to print the list of available languages, or"
		echo -e "     'q' to exit the program without creating an ISO: \c"
		read answer

		case ${answer} in
			p | P )
				printLanguages
				continue
				;;
			q | Q )
				return 1
				;;
			* )
				found=0
				for i in ${lTags[*]}; do
					if [[ ${answer} == $i ]]; then 
						found=1; 
						break
					fi;
				done

				if (( found == 0 )); then
				  echo -e "${answer} is not a valid language tag. Please try again"
				else
				  esdLang=${answer}
				fi
				;;
		esac
	done

	xpath -q -n -e '//File[LanguageCode="'${esdLang}'"]' "${wDir}/w11arm_products.xml" >"${wDir}/esd_edition.xml"
	esdURL=$(xpath -n -q -e '//FilePath' "${wDir}/esd_edition.xml" | sed -E -e 's/<[\/]?FilePath>//g')
	buildName=$(xpath -n -q -e '//FileName' "${wDir}/esd_edition.xml"| sed -E -e 's/(<FileName>)|(\.esd<[\/]FileName>)//g')
    shaHash=$(xpath -n -q -e '//Sha1' "${wDir}/esd_edition.xml" | sed -E -e 's/<[\/]?Sha1>//g')
	
	esdFile=${wDir}/${buildName}.esd
	isoFile=./${buildName}.iso

	echo -e "\nBuilding a Windows 11 ARM ${editionName} Editon ${esdLang} language ISO"
	echo -e "The generated ISO will be named ${isoFile}"
	if debugOn ; then
		echo
	    echo -e "[DEBUG] ESD Download variables"
	    echo -e "\n\tesdURL = ${esdURL}"
	    echo -e "\tbuildName = ${buildName}"
	    echo -e "\tesdFile = ${esdFile}"
	    echo -e "\tshaHash = ${shaHash}"
	    echo
	fi
	
	#---------------
	# Write out information that we'll need for restart
	# in a flag file restartOK that we set in the work directory
	#---------------

	echo -e "isoFile\t${isoFile}" >  "${wDir}/restartOK"
	echo -e "esdFile\t${esdFile}" >> "${wDir}/restartOK"
	echo -e "esdURL\t${esdURL}"   >> "${wDir}/restartOK"
	echo -e "shaHash\t${shaHash}" >> "${wDir}/restartOK"

	return 0
}

downloadEsd() {

	local retVal
	local -r retryLimit=10
	local displayInterval=0
	local calculatedHash
	local retryCount
	
	retryCount=0

	#---------------
	# Download the ESD from Microsoft
	#---------------
	
	echo
	verboseOn && displayInterval=120
	while (( retryCount < retryLimit )); do
		let retryCount+=1
		echo "Download attempt $retryCount"
		retVal=0
		aria2c --summary-interval=${displayInterval} --download-result=hide --dir="${wDir}" --file-allocation=none "${esdURL}" && break
		retVal=$?
		echo -e "Download interrupted"
	done
	(( retval != 0 )) && echo -n "Retry limit exceeded"
	return $retVal
}

extractEsd(){	
	
	local -r eFile=${esdFile}
	local -r eDir=${extDir}
	local retVal
	local esdImageCount
	local bootWimFile=${eDir}/sources/boot.wim
	local installWimFile=${eDir}/sources/install.wim

	local imageIndex
	local imageEdition
	local beQuiet="--quiet"
	
	
	verboseOn && beQuiet=""
	#---------------
	# Check the number of images in the esd. If 6, add it to the list of images in $images
	#---------------
		
	esdImageCount=$(wimlib-imagex info "${eFile}" | awk '/Image Count:/ {print $3}')
	verboseOn && echo -e "[NOTE] image count in ESD: $esdImageCount"

	#---------------
	# Extract image 1 in the ESD to create the boot environment
	#---------------

	echo -e "\nApplying boot files to the image"
	wimlib-imagex apply "$eFile" 1 "${eDir}" ${beQuiet} 2>/dev/null || {
		retVal=$?
		echo -e "[ERROR] Extract of boot files failed"
		return $retVal
	}

	echo -e "Boot files successfully applied to image"

	#---------------
	# Create the boot.wim file that contains WinPE and Windows Setup
	# Images 2 and 3 in the ESD contain these components
	#
	# Important: image 3 in the ESD must be marked as bootable when
	# transferred to boot.wim or else the installer will fail
	#---------------

	echo -e "\nAdding WinPE and Windows Setup to the image"
	wimlib-imagex export "${eFile}" 2 "${bootWimFile}" --compress=LZX --chunk-size 32K  ${beQuiet} || {
		retVal=$?
		echo -e "[ERROR] Add of WinPE failed"
		return ${retVal}
	}
	
	wimlib-imagex export "${eFile}" 3 "$bootWimFile" --compress=LZX --chunk-size 32K --boot  ${beQuiet} || {
		retVal=$?
		echo -e "[ERROR] Add of Windows Setup failed"
		return ${retVal}
	}
	echo -e "WinPE and Windows Setup added successfully to image\n"
	
	debugOn && {
		echo -e "[DEBUG] contents of ${bootWimFile}"
		wimlib-imagex info  "${bootWimFile}"
	}


	#---------------
	# Create the install.wim file that contains all images present in the ESD from
	# image index 4 onward - those will be available to Windows Setup
	#---------------
	
	echo -e "\nAdding Windows editions from the ESD to the image"
	for (( imageIndex=4; imageIndex<=esdImageCount; imageIndex++ )); do
		imageEdition=$(wimlib-imagex info "${eFile}" ${imageIndex} | grep '^Description:' | sed 's/Description:[ \t]*//')
		verboseOn && echo -e "\nAdding $imageEdition to the image"
		wimlib-imagex export "${eFile}" ${imageIndex} "${installWimFile}" --compress=LZMS --chunk-size 128K  ${beQuiet} || {
			retVal=$?
			echo -e "[ERROR] Addition of ${imageIndex} to the image failed"
			return $retVal
		}
		echo -e "${imageEdition} added successfully to the image"
	done

	echo -e "\nAll Windows editions added to image"	
	
	debugOn && {
		echo -e "[DEBUG] contents of ${installWimFile}"
		wimlib-imagex info "${installWimFile}"
	}
		
	return 0
}

buildIso(){
	local -r iDir=${extDir}
	local -r iFile=${isoFile}
	local -r bootFile=efi/microsoft/boot/efisys_noprompt.bin
	local sizeArg=""
	local beQuiet="-quiet"
	local retVal
	
	verboseOn && beQuiet=""
	
	if [[ -e "${iFile}" ]]; then
		echo -e "[INFO] File ${iFile} exists, removing it"
		rm -rf "${iFile}"
	fi
	
	# 
	# The non-Schily mkisofs needs an additional argument to properly build
	# Win 11 ARM 23H2 ISOs. See which mkisofs we have, and then add the -allow-limited-size
	# argument if we have a non Schily version
	#
	
	mkisofs -help 2>&1 | grep --quiet -- '-allow-limited-size'
	retVal=$?
	if (( retVal == 0 )) ; then
	    #
	    # Not Schily. Add -allow-limited-size argument
	    #
	    sizeArg="-allow-limited-size"
	    debugOn && echo "[DEBUG] Not Schily mkisofs"
	else
		debugOn && echo '[DEBUG] Schily mkisofs'
	fi

	mkisofs ${beQuiet} -b ${bootFile} -no-emul-boot -udf ${sizeArg} -iso-level 3 -hide "*" -V "ESD_ISO" -o "${iFile}" "$iDir"
	
	return $?
}


#-------------------
#
# Start of program
#
#-------------------

export LC_MESSAGES="C"

# Check to see if we're running on Darwin/macOS

if [[ $(uname -o) == "Darwin" ]]; then
	runningOnDarwin=1
fi


#-------------------
# 
# Process arguments (the -d option is undocumented and enables debug mode)
# 
#-------------------

while getopts ":hr:vVd" opt; do
  case ${opt} in
    h)
    	usage
    	exit 1
    	;;

    r)
    	rOption=1
    	wDir=$OPTARG
    	;;
    d)
    	dOption=1
    	let verbosityLevel+=1
    	;;
    	
	v)
    	let verbosityLevel+=1
    	;;
	V)
    	echo -e $version
    	exit 1
    	;;
    :)
    	echo -e "[ERROR] Option -$OPTARG requires an argument"
    	usage
    	exit 1
    	;;
    
    \?)
    	echo -e "[ERROR] Invalid option: -$OPTARG\n"
    	usage
    	exit 1
    	;;
    esac
done
shift "$((OPTIND-1))"


#-------------------
# Check number of arguments
# One argument is allowed when using the -r option for restart
# No arguments are allowed otherwise
#-------------------

if (( $# > 0 )); then
	echo -e "[ERROR] Too many arguments"
	usage
	exit 1
fi

verboseOn && echo -e ${version}
debugOn && echo -e "[DEBUG] Debug mode enabled"

#-------------------
# Check if required utilities are installed
#-------------------

notFound=""
for i in $requires; do
	which $i > /dev/null || notFound="${notFound} ${i}"
done
	
if [[ -n ${notFound} ]]; then
	abortProg 2 "Required utilities check" "\tRequired utilities are not found:\n\t${notFound}"
fi

if (( rOption == 0 )) ; then

	#---------------
	# Normal processing, no restart
	#---------------
		
	if (( runningOnDarwin == 1 )); then
		freeSpace=$(df -g "$tempDir" | awk '/^\/dev/ {print $4}' ) 
	else
		freeSpace=$(df -BG "$tempDir" | awk '/^\/dev/ {print $4}' | sed 's/G//' )
	fi

	verboseOn && echo -e "[NOTE] Free space in $tempDir is $freeSpace GB"
	if (( freeSpace < minDiskSpace )) ; then
		abortProg 28 "Disk space check" "\tThis utility requires $minDiskSpace GB of free disk space in the folder $tempDir\n\tand you have \n$freeSpace GB remaining."
	else
		if (( freeSpace < ( minDiskSpace + 3 ) )); then
			echo -e "[WARNING] This utility typically requires $minDiskSpace to $(( minDiskSpace+3 )) GB of free disk space"
			echo -e "          in the folder $tempDir"
			echo -e "          You have $freeSpace GB available. You may run out of disk space during the process."
		fi
	fi

    wDir=$(mktemp -q -d "${tempDir}/esd2iso_temp.XXXXXX")
	if (( $? != 0 )); then
		abortProg 2 "Work directory creation" "Unable to create work directory"
	fi
	verboseOn && echo "[INFO] Work directory $wDir created"
	
	extDir=${wDir}/ESD_ISO
	mkdir "${extDir}"

 	verboseOn && echo -e "\n[INFO] Windows product information is being downloaded from Microsoft"
	downloadCatalog "${wDir}" || {
		abortProg $? "Catalog download" ""
	}
	verboseOn && echo -e "[INFO]Product info download completed"
	
	echo -e "\nStep 1: Select the desired edition and language\n"
	
	setupEsdDownload || {
		rm -rf "${wDir}"
		echo -e "Progran exiting. No ISO was created"
		exit 0
	}
	
	echo -e "\nStep 1 complete"

	echo -e "\nStep 2: Downloading ESD from Microsoft"
	
else
	#---------------
	# Restart option selected
	#---------------
	
	#---------------
	# See if the work directory exists and if it supports restart
	#---------------
	
	[[ -d "${wDir}" ]] ||
		abortProg 2 "Restart" "\tWork directory ${wDir} does not exist"
	[[ -f "${wDir}/restartOK" ]] ||
		abortProg 2 "Restart" "\tThe work directory ${wDir} does not support restart"
	
	echo -e "\nRestarting interrupted download using work directory ${wDir}\n"

	#---------------
	#
	# Read variables needed for restart from the restartOK file
	#
	#---------------
	
	extDir=${wDir}/ESD_ISO
	isoFile=$(grep isoFile "${wDir}/restartOK" | cut -f 2)
	esdFile=$(grep esdFile "${wDir}/restartOK" | cut -f 2)
	esdURL=$(grep esdURL "${wDir}/restartOK" | cut -f 2)
	shaHash=$(grep shaHash "${wDir}/restartOK" | cut -f 2)
	
	debugOn && {
		echo -e "[DEBUG] Variables found in restart file:"
	    echo -e "\tesdURL = ${esdURL}"
	    echo -e "\tesdFile = ${esdFile}"
	    echo -e "\tisoFile = ${isoFile}"
	    echo -e "\tshaHash = ${shaHash}\n"
	}

	echo -e "Restarting Step 2: Downloading ESD from Microsoft"
fi

downloadEsd || 
	abortProg $? "ESD download" "\tRestart the download with the following command:\n$0 -r \"$wDir\""

rm "${wDir}/restartOK"

echo -e "\nDownload complete. Validating download (this may take a minute or so)"
calcHash=$(shasum "${esdFile}" | awk '{print $1}')
debugOn && {
	echo -e "[DEBUG] shaHash =  '${shaHash}'"
	echo -e "        calcHash = '${calcHash}'"
}
[[ ${calcHash} != ${shaHash} ]] &&
    abortProg 1 "Download validation" "\tExpected hash = ${shaHash}\n\tCalculated hash = ${calcHash}\n\tWork directory $wDir was not deleted, use for debugging"

echo -e "Download validated."
verboseOn && wimlib-imagex info "${esdFile}"


echo -e "\nStep 2 complete - ESD downloaded"	

echo -e "\nStep 3: Building installation image from ESD"

extractEsd || 
	abortProg $? "Image build" "Work directory ${wDir} was not deleted"

#---------------
# At this point we no longer need the ESD file as it's already extracted
# In order to reduce disk space reauirements, delete the ESD file unless we
# have the debug option set
#---------------

if debugOn ; then
	echo -e "[DEBUG] Keeping ESD download for debugging"
else
	echo -e "\nESD added successfully to installation image and is no longer needed.\nDeleting it to save disk space."
	verboseOn && echo -e "Deleting ESD file ${esdFile}"
	rm -rf "${esdFile}" 
	retVal=$?
	if (( retVal != 0 )); then
		echo -e "[WARNING] Deletion of ESD file encountered a problem."
		echo -e "          The ISO build can continue, but will consume an addtional 5 GB of disk space."
	else
		echo -e "ESD file deleted successfully\n"
	fi
fi

echo -e "\nStep 3 complete - installation image built"
	
echo -e "\nStep 4: Creating ISO $isoFile from the installation image\n"

buildIso  || 
    abortProg $? "ISO creation" "\tThe ISO was NOT created.\n\tWorking directory ${wDir} was not deleted"

echo
echo -e "Step 4 complete - ISO created"

if debugOn ; then
	echo -e "[DEBUG] Work directory $wDir was not deleted"
else
	echo -e "\nCleaning up work directory"
	rm -rf "${wDir}"
fi

echo -e "All done!"
exit 0
